from ollama import Client
import time
import os
import sys
from colorama import init, Fore, Style
import json
from datetime import datetime
import requests

def format_message(role, content):
    return {
        "role": role,
        "content": content,
        "created": int(datetime.now().timestamp())
    }

def send_chat_request(message, stream=True):
    url = "http://localhost:11434/api/chat"
    data = {
        "model": "qwen2-0.5b",
        "messages": [
            {"role": "user", "content": message}
        ],
        "stream": stream
    }
    
    response = requests.post(url, json=data, stream=stream)
    return response

def stream_chat():
    messages = []
    
    try:
        while True:
            user_input = input(f"{Fore.GREEN}User > {Style.RESET_ALL}")
            
            if user_input.lower() in ['exit', 'quit', '退出']:
                print(f"\n{Fore.GREEN}Session ended{Style.RESET_ALL}")
                break
            
            user_message = format_message("user", user_input)
            messages.append(user_message)
            
            print(f"{Fore.BLUE}Assistant > {Style.RESET_ALL}", end='')
            
            try:
                response = send_chat_request(user_input)
                
                if response.status_code == 200:
                    full_response = ""
                    for line in response.iter_lines():
                        if line:
                            chunk = json.loads(line.decode('utf-8'))
                            if 'message' in chunk and 'content' in chunk['message']:
                                content = chunk['message']['content']
                                sys.stdout.write(content)
                                sys.stdout.flush()
                                full_response += content
                    
                    if full_response:
                        assistant_message = format_message("assistant", full_response)
                        messages.append(assistant_message)
                    print("\n")
                else:
                    print_error(f"\nAPI error: {response.status_code}")
                    print(f"Response: {response.text}")
                
            except Exception as e:
                print_error(f"\nError: {str(e)}")
                print("Continuing to next round...\n")
                continue
                
    except ConnectionError as e:
        error_response = {
            "error": {
                "message": f"Connection error: {str(e)}\nPlease make sure 'ollama serve' is running",
                "type": "connection_error",
                "code": 500
            }
        }
        print(json.dumps(error_response, indent=2))
        
    except Exception as e:
        error_response = {
            "error": {
                "message": f"Error: {str(e)}",
                "type": "api_error",
                "code": 400
            }
        }
        print(json.dumps(error_response, indent=2))

def print_error(message):
    print(f"{Fore.RED}{message}{Style.RESET_ALL}")

def print_startup_info():
    startup_info = {
        "status": "online",
        "model": "qwen2-0.5b",
        "api_version": "v1",
        "server_time": datetime.now().isoformat()
    }
    print(json.dumps(startup_info, indent=2))

if __name__ == "__main__":
    init()
    os.system('')
    
    print_startup_info()
    # print("Waiting for Ollama service to initialize...")
    # time.sleep(1)
    
    stream_chat()